import torchvision.transforms
import config
import torch
import torchvision.transforms as trans
from torch.utils.data import Dataset  # 导入PyTorch数据集基类
import pandas as pd  # 用于读取CSV文件
from PIL import Image  # 用于图像处理
import os  # 用于路径操作


class YoloV1Dataset (Dataset) :
    def __init__ (
            self, data_csv, img_dir, label_dir,
            S=config.S, B=config.B, C=config.C, img_size=config.IMAGE_SIZE,
            normalize=False
    ) :
        """
        初始化YOLOv1数据集类
        :param data_csv: 包含图像路径和标签路径的CSV文件路径（列格式：[img_path, label_path]）
        :param img_dir: 图像文件所在目录
        :param label_dir: 标签文件所在目录
        :param S: 图像划分的网格数（S×S），对应config中的S
        :param B: 每个网格预测的边界框数量，对应config中的B
        :param C: 目标类别数，对应config中的C
        :param img_size: 输入网络的图像尺寸（宽×高），默认使用config中的IMAGE_SIZE
        :param normalize: 是否对图像进行归一化处理（默认不启用）
        """
        self.df = pd.read_csv (data_csv)  # 读取包含图像和标签路径的CSV文件
        self.img_dir = img_dir  # 图像目录路径
        self.label_dir = label_dir  # 标签目录路径
        self.S = S  # 网格数量S
        self.B = B  # 每个网格的边界框数B
        self.C = C  # 类别数C
        self.depth = self.B * 5 + self.C  # 输出张量的深度（2个边界框×5属性 + 20类别）
        self.img_size = img_size  # 输入图像尺寸
        self.normalize = normalize  # 归一化标志位

    def __len__ (self) :
        """
        返回数据集的样本数量
        :return: 数据集长度（整数）
        """
        return len (self.df)  # 返回CSV文件中的行数（即样本总数）

    def __getitem__ (self, idx) :
        """
        根据索引获取单个样本（图像+标签）
        :param idx: 样本索引（0到len-1）
        :return: 图像数据（Tensor）和标签数据（Tensor）
        """
        # 获取图像和标签的完整路径
        img_path = os.path.join (self.img_dir, self.df.iloc[idx, 0])  # 拼接图像路径（目录+文件名）
        label_path = os.path.join (self.label_dir, self.df.iloc[idx, 1])  # 拼接标签路径（目录+文件名）

        # 读取图像和标签
        img = Image.open (img_path)  # 使用PIL打开图像文件
        label = self.read_label (label_path)  # 调用read_label方法解析标签文件

        # 图像预处理（调整尺寸+转换为Tensor+可选归一化）
        if self.normalize :  # 启用归一化时
            transforms = trans.Compose ([
                trans.Resize (self.img_size),  # 调整图像尺寸为(448,448)
                trans.ToTensor (),  # 将PIL图像转换为Tensor（范围0-1）
                trans.Normalize (
                    [0.485, 0.456, 0.406],  # 均值（ImageNet统计值）
                    [0.22, 0.224, 0.225]     # 标准差（ImageNet统计值）
                )
            ])
            img = transforms (img)  # 应用预处理
        else :  # 不启用归一化时
            transforms = trans.Compose ([trans.Resize (self.img_size), trans.ToTensor ()])  # 仅调整尺寸和转换Tensor
            img = transforms (img)  # 应用预处理

        return img, label  # 返回预处理后的图像和标签张量


    def read_label (self, label_path) :
        """
        读取标签文件并转换为模型所需的张量格式
        :param label_path: 标签文件路径（文本文件）
        :return: 标签矩阵（形状：(S, S, depth)，默认(7,7,30)）
        """
        label_matrix = torch.zeros((self.S, self.S, self.depth))  # 初始化全零标签矩阵（S×S网格，每个网格depth维）
        labels = []  # 存储原始标签列表（每行一个边界框信息）
        with open(label_path, 'r') as f :  # 打开标签文件
            while True :
                line = f.readline()
                if not line :  # 读取到空行时结束
                    break
                labels.append(line.strip().split())  # 按空格分割每行数据（格式：[类别, x中心, y中心, 宽, 高]）
                # 最终labels格式：[[l, x, y, w, h], [l, x, y, w, h], ...]

        for idx, label in enumerate (labels) :  # 遍历每个边界框
            # 计算边界框中心所在的网格坐标（i,j）
            grid_i = int (self.S * float (label[1]))  # 水平方向网格索引（0~S-1）
            grid_j = int (self.S * float (label[2]))  # 垂直方向网格索引（0~S-1）

            # 计算边界框相对于网格的偏移量和尺寸（归一化到0-1）
            cell_x = self.S * float (label[1]) - grid_i  # 中心x坐标在网格内的相对位置（0~1）
            cell_y = self.S * float (label[2]) - grid_j  # 中心y坐标在网格内的相对位置（0~1）
            cell_w = self.S * float (label[3])  # 边界框宽度相对于图像的比例（0~1）
            cell_h = self.S * float (label[4])  # 边界框高度相对于图像的比例（0~1）

            # 将边界框信息填充到标签矩阵中
            label_matrix[grid_i][grid_j][0] = 1                     # 第一个边界框的置信度（存在目标为1）
            label_matrix[grid_i][grid_j][1] = cell_x                # 相对网格的x偏移
            label_matrix[grid_i][grid_j][2] = cell_y                # 相对网格的y偏移
            label_matrix[grid_i][grid_j][3] = cell_w                # 相对图像的宽度比例
            label_matrix[grid_i][grid_j][4] = cell_h                # 相对图像的高度比例
            label_matrix[grid_i][grid_j][int (label[0]) + 10] = 1   # 类别概率（对应类别位置设为1）

        return label_matrix  # 返回构造好的标签矩阵




if __name__ == '__main__' :
    # img = Image.open ("./datasets/images/000001.jpg")
    # transforms = trans.Compose ([trans.ToTensor ()])
    # print (transforms (img))
    # transforms = trans.Compose ([trans.ToTensor (), trans.Normalize([0.485, 0.456, 0.405], [0.2929, 0.224, 0.225])])
    # print (transforms (img))

    # dataset = YoloV1Dataset ("./datasets/train.csv", "./datasets/images", "./datasets/labels")
    # img, label = dataset.__getitem__(0)
    # print (img.shape)
    # print (label)

    # t = torch.zeros((3, 4, 2))
    # t[0, 1] = 2.5
    # print (t)

    train_data = YoloV1Dataset ("./datasets/train.csv", "./datasets/images", "./datasets/labels", normalize=True)
    print (train_data.__len__())
    print (train_data.__getitem__(1))